# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Bbox utils"""

import math
import itertools as it
import numpy as np
from src.model_utils.config import config

class GeneratDefaultBoxes():
    """
    Generate Default boxes for nanodet, follows the order of (W, H, archor_sizes).
    `self.default_boxes` has a shape of [archor_sizes, H, W, 4], the last dimension is [y, x, h, w].
    `self.default_boxes_ltrb` has a shape as `self.default_boxes`, the last dimension is [y1, x1, y2, x2].
    """
    def __init__(self):
        octave_base_scale = config.octave_base_scale
        scales = np.array([2 ** 0])
        anchor_size = np.array(config.anchor_size)
        self.default_boxes = []
        for idex, feature_size in enumerate(config.feature_size):
            base_size = anchor_size[idex]
            size1 = base_size*scales[0]
            all_sizes = []
            for aspect_ratio in config.aspect_ratios[idex]:
                w1, h1 = size1 * math.sqrt(aspect_ratio), size1 / math.sqrt(aspect_ratio)
                all_sizes.append((h1, w1))
            for i, j in it.product(range(feature_size), repeat=2):
                for h, w in all_sizes:
                    cx, cy = (j + 0.5) * config.steps[idex], (i + 0.5) * config.steps[idex]
                    self.default_boxes.append([cx, cy, h, w])

        def to_ltrb(cx, cy, h, w):
            h, w = h * octave_base_scale, w * octave_base_scale
            return cx - h / 2, cy - w / 2, cx + h / 2, cy + w / 2

        # For IoU calculation
        self.default_boxes_ltrb = np.array(tuple(to_ltrb(*i) for i in self.default_boxes), dtype='float32')
        self.default_boxes = np.array(self.default_boxes, dtype='float32')


default_boxes_ltrb = GeneratDefaultBoxes().default_boxes_ltrb
default_boxes = GeneratDefaultBoxes().default_boxes
x1, y1, x2, y2 = np.split(default_boxes_ltrb[:, :4], 4, axis=-1)
vol_anchors = (x2 - x1) * (y2 - y1)
num_level_cells_list = [1600, 400, 100]


def nanodet_bboxes_encode(boxes, img_id):
    """
    Labels anchors with ground truth inputs.

    Args:
        boxex: ground truth with shape [N, 5], for each row, it stores [y, x, h, w, cls].

    Returns:
        gt_loc: location ground truth with shape [num_anchors, 4].
        gt_label: class ground truth with shape [num_anchors, 1].
        num_matched_boxes: number of positives in an image.
    """

    def jaccard_with_anchors(bbox):
        """Compute jaccard score a box and the anchors."""
        # Intersection bbox and volume.
        xmin = np.maximum(x1, bbox[0])
        ymin = np.maximum(y1, bbox[1])
        xmax = np.minimum(x2, bbox[2])
        ymax = np.minimum(y2, bbox[3])
        w = np.maximum(ymax - ymin, 0.)
        h = np.maximum(xmax - xmin, 0.)

        # Volumes.
        inter_vol = h * w
        union_vol = vol_anchors + (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) - inter_vol
        jaccard = inter_vol / union_vol
        return np.squeeze(jaccard)

    pre_scores = np.zeros((config.num_nanodet_boxes), dtype=np.float32)
    t_boxes = np.zeros((config.num_nanodet_boxes, 4), dtype=np.float32)
    t_label = np.full((config.num_nanodet_boxes), 80, dtype=np.int64)

    candidate_idxs_list = []
    scores_list = []
    labels = []
    num_gt = boxes.shape[0]
    # ATSS
    for bbox in boxes:
        label = int(bbox[4])
        labels.append(label)
        scores = jaccard_with_anchors(bbox)
        scores_list.append(scores)
        gt_cx = (bbox[0] + bbox[2]) / 2.0
        gt_cy = (bbox[1] + bbox[3]) / 2.0
        gt_points = np.stack((gt_cx, gt_cy), axis=0)
        grid_priors_cx = (default_boxes_ltrb[:, 0] + default_boxes_ltrb[:, 2]) / 2.0
        grid_priors_cy = (default_boxes_ltrb[:, 1] + default_boxes_ltrb[:, 3]) / 2.0
        grid_cells_points = np.stack((grid_priors_cx, grid_priors_cy), axis=1)
        distances = grid_cells_points - gt_points
        distances = np.power(distances, 2)
        distances = np.sum(distances, axis=-1)
        distances = np.sqrt(distances)
        candidate_idxs = []
        start_idx = 0
        topk = 9
        for level, cells_per_level in enumerate(num_level_cells_list):
            end_idx = start_idx + cells_per_level
            distances_per_level = distances[start_idx:end_idx]
            selectable_k = min(topk, cells_per_level)
            topk_idxs_per_level = np.argsort(distances_per_level, axis=0)
            topk_idxs_per_level = topk_idxs_per_level[:selectable_k]
            candidate_idxs.append(topk_idxs_per_level + start_idx)
            start_idx = end_idx
        candidate_idxs = np.concatenate(candidate_idxs, axis=0)
        candidate_idxs_list.append(candidate_idxs)
        
    candidate_idxs = np.stack(candidate_idxs_list, axis=1)
    scores = np.stack(scores_list, axis=1)
    labels = np.stack(labels, axis=0)

    candidate_overlaps = scores[candidate_idxs, np.arange(num_gt)]
    overlaps_mean = candidate_overlaps.mean(0)
    overlaps_std = candidate_overlaps.std(0)
    overlaps_thr = overlaps_mean + overlaps_std
    is_pos = (candidate_overlaps >= overlaps_thr)
    l_ = grid_priors_cx[candidate_idxs] - boxes[:, 0]
    t_ = grid_priors_cy[candidate_idxs] - boxes[:, 1]
    r_ = boxes[:, 2] - grid_priors_cx[candidate_idxs]
    b_ = boxes[:, 3] - grid_priors_cy[candidate_idxs]
    is_in_gts = np.min(np.stack([l_, t_, r_, b_], axis=0), axis=0) > 0.01
    is_pos = is_pos & is_in_gts
    max_score_indx = np.argmax(scores,1)
    idx = candidate_idxs[is_pos]

    for i in range(num_gt):
        idx = candidate_idxs[:, i][is_pos[:, i]]
        score = scores[:, i]
        mask = np.zeros((config.num_nanodet_boxes,), np.bool)
        mask[idx] = True
        mask = mask & (score > pre_scores)
        pre_scores = np.maximum(pre_scores, score * mask)
        t_label = mask * boxes[i][4] + (1 - mask) * t_label
        for c in range(4):
            t_boxes[:, c] = mask * boxes[i][c] + (1 - mask) * t_boxes[:, c]
    is_pos_label = t_label < 80
    index = np.nonzero(is_pos_label)
    num_match = np.array([np.maximum(len(np.nonzero(is_pos_label)[0]), 1)], dtype=np.int32)
    return t_boxes, t_label.astype(np.int32), default_boxes, num_match

def intersect(box_a, box_b):
    """Compute the intersect of two sets of boxes."""
    max_yx = np.minimum(box_a[:, 2:4], box_b[2:4])
    min_yx = np.maximum(box_a[:, :2], box_b[:2])
    inter = np.clip((max_yx - min_yx), a_min=0, a_max=np.inf)
    return inter[:, 0] * inter[:, 1]


def jaccard_numpy(box_a, box_b):
    """Compute the jaccard overlap of two sets of boxes."""
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2] - box_a[:, 0]) *
              (box_a[:, 3] - box_a[:, 1]))
    area_b = ((box_b[2] - box_b[0]) *
              (box_b[3] - box_b[1]))
    union = area_a + area_b - inter
    return inter / union
